// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <kthread/internal/types.h>
#include <tally/tally.h>
#include <turbo/utility/errno.h>

__BEGIN_DECLS
extern int kthread_mutex_init(kthread_mutex_t *__restrict mutex,
                              const kthread_mutexattr_t *__restrict mutex_attr);
extern int kthread_mutex_destroy(kthread_mutex_t *mutex);
extern int kthread_mutex_trylock(kthread_mutex_t *mutex);
extern int kthread_mutex_lock(kthread_mutex_t *mutex);
extern int kthread_mutex_timedlock(kthread_mutex_t *__restrict mutex,
                                   const struct timespec *__restrict abstime);
extern int kthread_mutex_unlock(kthread_mutex_t *mutex);
__END_DECLS

namespace kthread {

// The C++ Wrapper of kthread_mutex

// NOTE: Not aligned to cacheline as the container of Mutex is practically aligned
    class Mutex {
    public:
        typedef kthread_mutex_t *native_handler_type;

        Mutex() {
            int ec = kthread_mutex_init(&_mutex, nullptr);
            if (ec != 0) {
                throw std::system_error(std::error_code(ec, std::system_category()), "Mutex constructor failed");
            }
        }

        ~Mutex() { KCHECK_EQ(0, kthread_mutex_destroy(&_mutex)); }

        native_handler_type native_handler() { return &_mutex; }

        void lock() {
            int ec = kthread_mutex_lock(&_mutex);
            if (ec != 0) {
                throw std::system_error(std::error_code(ec, std::system_category()), "Mutex lock failed");
            }
        }

        void unlock() { kthread_mutex_unlock(&_mutex); }

        bool try_lock() { return !kthread_mutex_trylock(&_mutex); }
        // TODO(chenzhangyi01): Complement interfaces for C++11
    private:
        TURBO_DISALLOW_COPY_AND_ASSIGN(Mutex);

        kthread_mutex_t _mutex;
    };

    namespace internal {
        class FastPthreadMutex {
        public:
            FastPthreadMutex() : _futex(0) {}

            ~FastPthreadMutex() = default;

            void lock();

            void unlock();

            bool try_lock();

        private:
            TURBO_DISALLOW_COPY_AND_ASSIGN(FastPthreadMutex);

            int lock_contended();

            unsigned _futex;
        };
    }

    class FastPthreadMutex {
    public:
        FastPthreadMutex() = default;

        ~FastPthreadMutex() = default;
        TURBO_DISALLOW_COPY_AND_ASSIGN(FastPthreadMutex);

        void lock();

        void unlock();

        bool try_lock() { return _mutex.try_lock(); }

    private:
        internal::FastPthreadMutex _mutex;
    };

    bool ContentionProfilerStart(const char *filename);
    void ContentionProfilerStop();

}  // namespace kthread

// Specialize std::lock_guard and std::unique_lock for kthread_mutex_t

namespace std {

    template<>
    class lock_guard<kthread_mutex_t> {
    public:
        explicit lock_guard(kthread_mutex_t &mutex) : _pmutex(&mutex) {
#if !defined(NDEBUG)
            const int rc = kthread_mutex_lock(_pmutex);
            if (rc) {
                KLOG(FATAL) << "Fail to lock kthread_mutex_t=" << _pmutex << ", " << km_error(rc);
                _pmutex = nullptr;
            }
#else
            kthread_mutex_lock(_pmutex);
#endif  // NDEBUG
        }

        ~lock_guard() {
#ifndef NDEBUG
            if (_pmutex) {
                kthread_mutex_unlock(_pmutex);
            }
#else
            kthread_mutex_unlock(_pmutex);
#endif
        }

    private:
        TURBO_DISALLOW_COPY_AND_ASSIGN(lock_guard);

        kthread_mutex_t *_pmutex;
    };

    template<>
    class unique_lock<kthread_mutex_t> {
        TURBO_DISALLOW_COPY_AND_ASSIGN(unique_lock);

    public:
        typedef kthread_mutex_t mutex_type;

        unique_lock() : _mutex(nullptr), _owns_lock(false) {}

        explicit unique_lock(mutex_type &mutex)
                : _mutex(&mutex), _owns_lock(false) {
            lock();
        }

        unique_lock(mutex_type &mutex, defer_lock_t)
                : _mutex(&mutex), _owns_lock(false) {}

        unique_lock(mutex_type &mutex, try_to_lock_t)
                : _mutex(&mutex), _owns_lock(kthread_mutex_trylock(&mutex) == 0) {}

        unique_lock(mutex_type &mutex, adopt_lock_t)
                : _mutex(&mutex), _owns_lock(true) {}

        ~unique_lock() {
            if (_owns_lock) {
                unlock();
            }
        }

        void lock() {
            if (!_mutex) {
                KCHECK(false) << "Invalid operation";
                return;
            }
            if (_owns_lock) {
                KCHECK(false) << "Detected deadlock issue";
                return;
            }
            kthread_mutex_lock(_mutex);
            _owns_lock = true;
        }

        bool try_lock() {
            if (!_mutex) {
                KCHECK(false) << "Invalid operation";
                return false;
            }
            if (_owns_lock) {
                KCHECK(false) << "Detected deadlock issue";
                return false;
            }
            _owns_lock = !kthread_mutex_trylock(_mutex);
            return _owns_lock;
        }

        void unlock() {
            if (!_owns_lock) {
                KCHECK(false) << "Invalid operation";
                return;
            }
            if (_mutex) {
                kthread_mutex_unlock(_mutex);
                _owns_lock = false;
            }
        }

        void swap(unique_lock &rhs) {
            std::swap(_mutex, rhs._mutex);
            std::swap(_owns_lock, rhs._owns_lock);
        }

        mutex_type *release() {
            mutex_type *saved_mutex = _mutex;
            _mutex = nullptr;
            _owns_lock = false;
            return saved_mutex;
        }

        mutex_type *mutex() { return _mutex; }

        bool owns_lock() const { return _owns_lock; }

        operator bool() const { return owns_lock(); }

    private:
        mutex_type *_mutex;
        bool _owns_lock;
    };

}  // namespace std

namespace tally {

    template<>
    struct MutexConstructor<kthread_mutex_t> {
        bool operator()(kthread_mutex_t *mutex) const {
            return kthread_mutex_init(mutex, nullptr) == 0;
        }
    };

    template<>
    struct MutexDestructor<kthread_mutex_t> {
        bool operator()(kthread_mutex_t *mutex) const {
            return kthread_mutex_destroy(mutex) == 0;
        }
    };

}  // namespace tally
